﻿/* Copyright 2007-2008 dnAnalytics Project.
 *
 * Contributors to this file:
 * Jurgen Van Gael
 *
 * Redistribution and use in source and binary forms, with or without modification,
 * are permitted provided that the following conditions are met:
 * 
 * * Redistributions of source code must retain the above copyright notice, this 
 *   list of conditions and the following disclaimer.
 * * Redistributions in binary form must reproduce the above copyright notice, 
 *   this list of conditions and the following disclaimer in the documentation
 *   and/or other materials provided with the distribution.
 * * Neither the name of the dnAnalytics Project nor the names of its contributors
 *   may be used to endorse or promote products derived from this software without
 *   specific prior written permission.
 * 
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
 * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
 * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE 
 * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
 * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
 * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
 * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
 * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR
 * TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF
 * THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 */

using System;
using System.Collections.Generic;
using dnAnalytics.LinearAlgebra;
using dnAnalytics.LinearAlgebra.Solvers.Direct;
using dnAnalytics.Properties;

namespace dnAnalytics.Statistics.Distributions
{
    /// <summary>
    /// The Conway-Maxwell-Poisson distribution is a generalization of the Poisson, Geometric and Bernoulli
    /// distributions. It is parameterized by two real numbers "lambda" and "nu". For
    ///     nu = 0 the distribution reverts to a Geometric distribution,
    ///     nu = 1 the distribution reverts to the Poisson distribution,
    ///     nu -> infinity the distribution converges to a Bernoulli distribution.
    /// 
    /// This implementation will cache the value of the normalization constant.
    /// </summary>
    /// <remarks>The distribution will use the <see cref="System.Random"/> by default. 
    /// Users can set the random number generator by using the <see cref="RandomNumberGenerator"/> property.
    /// 
    /// The statistics classes will check all the incoming parameters whether they are in the allowed
    /// range. This might involve heavy computation. Optionally, by setting Control.CheckDistributionParameters
    /// to false, all parameter checks can be turned off.</remarks>
    public class ConwayMaxwellPoisson : IDiscreteDistribution
    {
        /// <summary>
        /// Since many properties of the distribution can only be computed approximately, the tolerance
        /// level specifies how much error we accept.
        /// </summary>
        private const double mTolerance = 1e-12;

        private readonly double mLambda;
        private readonly double mNu;

        private double mMean = double.MinValue;
        private double mVariance = double.MinValue;

        /// <summary>
        /// Caches the value of the normalization constant.
        /// </summary>
        private double mZ = double.MinValue;

        /// <summary>
        /// Build a new Conway-Maxwell-Poisson distribution.
        /// </summary>
        /// <param name="lambda">The lambda parameter.</param>
        /// <param name="nu">The nu parameter.</param>
        public ConwayMaxwellPoisson(double lambda, double nu)
        {
            if (Control.CheckDistributionParameters)
            {
                CheckParameters(nu, lambda);
            }

            mNu = nu;
            mLambda = lambda;
            RandomNumberGenerator = new System.Random();
        }

        private static void CheckParameters(double lambda, double nu)
        {
            if (lambda <= 0.0)
            {
                throw new ArgumentOutOfRangeException("lambda", Resources.NotPositive);
            }
            if (nu < 0.0)
            {
                throw new ArgumentOutOfRangeException("nu", Resources.ParameterCannotBeNegative);
            }
        }

        /// <summary>
        /// Gets the lambda parameter.
        /// </summary>
        /// <value>The value of the lambda parameter.</value>
        public double Lambda
        {
            get { return mLambda; }
        }

        /// <summary>
        /// Gets the Nu parameter.
        /// </summary>
        /// <value>The value of the Nu parameter.</value>
        public double Nu
        {
            get { return mNu; }
        }

        /// <summary>
        /// A string representation of the distribution.
        /// </summary>
        public override string ToString()
        {
            return "ConwayMaxwellPoisson(Lambda = " + mLambda + ", Nu = " + mNu + ")";
        }

        /// <summary>
        /// The normalization constant of the Conway-Maxwell-Poisson distribution.
        /// </summary>
        public double Z
        {
            get
            {
                if (mZ != double.MinValue)
                {
                    return mZ;
                }
                mZ = Normalization(mLambda, mNu);
                return mZ;
            }
        }

        #region IDistribution

        /// <summary>
        /// The mean of the distribution.
        /// </summary>
        /// <value></value>
        public double Mean
        {
            get
            {
                // Special case requiring no computation.
                if (mLambda == 0)
                {
                    return 0.0;
                }
                if (mMean != double.MinValue)
                {
                    return mMean;
                }
                // The normalization constant for the distribution.
                double z = 1 + mLambda;
                // The probability of the next term.
                double a1 = mLambda*mLambda/System.Math.Pow(2, mNu);
                // The unnormalized mean.
                double zx = mLambda;
                // The contribution of the next term to the mean.
                double ax1 = 2*a1;

                for (int i = 3; i < 1000; i++)
                {
                    double e = mLambda/System.Math.Pow(i, mNu);
                    double ex = mLambda/System.Math.Pow(i, mNu - 1)/(i - 1);
                    double a2 = a1*e;
                    double ax2 = ax1*ex;

                    double m = zx/z;
                    double upper = (zx + ax1/(1 - ax2/ax1))/z;
                    double lower = zx/(z + a1/(1 - a2/a1));

                    if ((ax2 < ax1) && (a2 < a1))
                    {
                        double r = (upper - lower)/m;
                        if (r < mTolerance)
                        {
                            break;
                        }
                    }

                    z = z + a1;
                    zx = zx + ax1;
                    a1 = a2;
                    ax1 = ax2;
                }

                mMean = zx/z;
                return mMean;
            }
        }

        /// <summary>
        /// The standard deviation of the distribution.
        /// </summary>
        /// <value></value>
        public double StdDev
        {
            get { return System.Math.Sqrt(Variance); }
        }

        /// <summary>
        /// The variance of the distribution.
        /// </summary>
        /// <value></value>
        public double Variance
        {
            get
            {
                // Special case requiring no computation.
                if (mLambda == 0)
                {
                    return 0.0;
                }

                if (mVariance != double.MinValue)
                {
                    return mVariance;
                }

                // The normalization constant for the distribution.
                double z = 1 + mLambda;
                // The probability of the next term.
                double a1 = mLambda*mLambda/System.Math.Pow(2, mNu);
                // The unnormalized second moment.
                double zxx = mLambda;
                // The contribution of the next term to the second moment.
                double axx1 = 4*a1;

                for (int i = 3; i < 1000; i++)
                {
                    double e = mLambda/System.Math.Pow(i, mNu);
                    double exx = mLambda/System.Math.Pow(i, mNu - 2)/(i - 1)/(i - 1);
                    double a2 = a1*e;
                    double axx2 = axx1*exx;

                    double m = zxx/z;
                    double upper = (zxx + axx1/(1 - axx2/axx1))/z;
                    double lower = zxx/(z + a1/(1 - a2/a1));

                    if ((axx2 < axx1) && (a2 < a1))
                    {
                        double r = (upper - lower)/m;
                        if (r < mTolerance)
                        {
                            break;
                        }
                    }

                    z = z + a1;
                    zxx = zxx + axx1;
                    a1 = a2;
                    axx1 = axx2;
                }

                double mean = Mean;
                mVariance = zxx/z - mean*mean;
                return mVariance;
            }
        }

        /// <summary>
        /// The entropy of the distribution.
        /// </summary>
        /// <value></value>
        public double Entropy
        {
            get { throw new Exception("Not implemented yet."); }
        }

        /// <summary>
        /// Gets or sets the random number generator.
        /// </summary>
        /// <value>The random number generator used to generate a random sample.</value>
        public System.Random RandomNumberGenerator { get; set; }

        #endregion

        #region IDiscreteDistribution Members

        /// <summary>
        /// The median of the distribution.
        /// </summary>
        /// <value></value>
        public int Median
        {
            get { throw new Exception("Not implemented yet"); }
        }

        /// <summary>
        /// The mode of the distribution.
        /// </summary>
        /// <value></value>
        public int Mode
        {
            get { throw new Exception("Not implemented yet"); }
        }

        /// <summary>
        /// Computes values of the probability mass function.
        /// </summary>
        /// <param name="k">The location in the domain where we want to evaluate the probability mass function.</param>
        /// <returns></returns>
        public double Probability(int k)
        {
            return System.Math.Pow(mLambda, k) / System.Math.Pow(Math.SpecialFunctions.Factorial(k), mNu) / Z;
        }

        /// <summary>
        /// Samples a Conway-Maxwell-Poisson distributed random variable.
        /// </summary>
        public int Sample()
        {
            return DoSample(RandomNumberGenerator, mLambda, mNu, mZ);
        }

        /// <summary>
        /// Samples an array of Conway-Maxwell-Poisson distributed random variables.
        /// </summary>
        /// <param name="n">The number of variables needed.</param>
        public int[] Sample(int n)
        {
            return DoSample(RandomNumberGenerator, n, mLambda, mNu, mZ);
        }

        #endregion

        /// <summary>
        /// Samples a Conway-Maxwell-Poisson distributed random variable.
        /// </summary>
        /// <param name="rnd">The random number generator to use.</param>
        /// <param name="lambda">The lambda parameter for the CMP distribution.</param>
        /// <param name="nu">The nu parameter for the CMP distribution.</param>
        public static int Sample(System.Random rnd, double lambda, double nu)
        {
            if (Control.CheckDistributionParameters)
            {
                CheckParameters(lambda, nu);
            }

            double Z = Normalization(lambda, nu);
            return DoSample(rnd, lambda, nu, Z);
        }

        private static int DoSample(System.Random rnd, double lambda, double nu, double z)
        {
            double u = rnd.NextDouble();
            double P = 1.0/z;
            double cdf = P;
            int i = 0;

            while (u > cdf)
            {
                i++;
                P = P*lambda/System.Math.Pow(i, nu);
                cdf += P;
            }

            return i;
        }

        private static int[] DoSample(System.Random rnd, int n, double lambda, double nu, double z)
        {
            int[] arr = new int[n];

            for (int i = 0; i < n; i++)
            {
                arr[i] = DoSample(rnd, lambda, nu, z);
            }

            return arr;
        }

        /// <summary>
        /// Samples an array of Conway-Maxwell-Poisson distributed random variables.
        /// </summary>
        /// <param name="rnd">The random number generator to use.</param>
        /// <param name="n">The number of variables needed.</param>
        /// <param name="lambda">The lambda parameter for the CMP distribution.</param>
        /// <param name="nu">The nu parameter for the CMP distribution.</param>
        public static int[] Sample(System.Random rnd, int n, double lambda, double nu)
        {
            if (Control.CheckDistributionParameters)
            {
                CheckParameters(lambda, nu);
            }

            double Z = Normalization(lambda, nu);
            return DoSample(rnd, n, lambda, nu, Z);
        }

        /// <summary>
        /// Computes an approximate normalization constant for the CMP distribution.
        /// </summary>
        /// <param name="lambda">The lambda parameter for the CMP distribution.</param>
        /// <param name="nu">The nu parameter for the CMP distribution.</param>
        public static double Normalization(double lambda, double nu)
        {
            // Initialize Z with the first two terms.
            double Z = 1.0 + lambda;

            // Remembers the last term added.
            double t = lambda;

            // Start adding more terms until convergence.
            for (int i = 2; i < 1000; i++)
            {
                // The new addition for term i.
                double e = lambda/System.Math.Pow(i, nu);
                // The new term.
                t = t*e;
                // The updated normalization constant.
                Z = Z + t;

                // The stopping criterion.
                if (e < 1)
                {
                    if (t/(1 - e)/Z < mTolerance)
                    {
                        break;
                    }
                }
            }

            return Z;
        }

        /// <summary>
        /// Estimate a new Conway-Maxwell-Poisson distribution by performing a regression on the log probability ratios. The inputs
        /// to this method should be sorted by ascending x.
        /// </summary>
        /// <param name="x">The indices for the counts.</param>
        /// <param name="counts">An array of counts.</param>
        /// <returns>A Conway-Maxwell-Poisson distribution with estimated parameters.</returns>
        public static ConwayMaxwellPoisson FromLogRatios(int[] x, double[] counts)
        {
            int n = x.Length;

            if (n != counts.Length)
            {
                throw new ArgumentException(Resources.ParametersNotConformable);
            }

            // A list to store the log probability ratios.
            List<double> logpr = new List<double>();
            // A list to store the indices.
            List<double> logx = new List<double>();

            // Make a list of indices and probabilities when the probability ratios are not going to involve 0's.
            for (int i = 1; i < n; i++)
            {
                if (counts[i] != 0.0 && counts[i - 1] != 0.0 && x[i] == x[i - 1] + 1)
                {
                    logpr.Add(System.Math.Log(counts[i - 1]/counts[i]));
                    logx.Add(System.Math.Log(x[i]));
                }
            }

            // The number of points to be used in the regression.
            int N = logpr.Count;

            // Setup the least squares problem.
            Vector y = new DenseVector(logpr.ToArray());
            Matrix X = new DenseMatrix(N, 2);
            for (int i = 0; i < N; i++)
            {
                X[i, 0] = 1.0;
                X[i, 1] = logx[i];
            }

            // Solve the least squares problem.
            CholeskySolver solver = new CholeskySolver();
            Matrix Xt = X.Transpose();
            Vector b = solver.Solve(Xt * X, Xt * y);

            // Extract the parameters for the CMP distribution.
            return new ConwayMaxwellPoisson(System.Math.Exp(-b[0]), b[1]);
        }
    }
}